import cv2
import os, sys
import glob
import numpy as np
import time
import dlib
import argparse
from pathlib2 import Path

detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor('../data/face/shape_predictor_68_face_landmarks.dat')

def rect_to_bbox(rect):
    x = rect.left()
    y = rect.top()
    w = rect.right() - x
    h = rect.bottom() - y

    return (x, y, w, h)

def shape_to_np(landmk, dtype="int"):
    # initialize the list of (x, y)-coordinates
    coords = np.zeros((landmk.num_parts, 2), dtype=dtype)

    # loop over all facial landmarks and convert them
    # to a 2-tuple of (x, y)-coordinates
    for i in range(0, landmk.num_parts):
        coords[i] = (landmk.part(i).x, landmk.part(i).y)

    # return the list of (x, y)-coordinates
    return coords

def detect_image(path_image, path_landmk, path_landmk_5p, path_landmk_bbox, image=None, is_save=True):

    # path_save = path_image.parent.joinpath(path_image.stem + ".txt")
    if not path_image.exists():
        raise(f"ERROR: {path_image} not exist")

    if image is None:
        image = cv2.imread(str(path_image))

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    rects = detector(gray, 1)

    if len(rects) == 0:
        return False
    elif len(rects) > 1:  # maybe have multi person
        print(f"WARRNING: have {len(rects)} face in this images!!!! select the first one!!!")
    
    for (i, rect) in enumerate(rects):
        if i>0: break    # only detect the first person, TODO: expand multiperson

        landmk = predictor(gray, rect)
        landmk = shape_to_np(landmk)
        landmk_bbox = rect_to_bbox(rect)

        #print(f"rect type = {type(rect)}, value={rect}, (x,y,w,h)={landmk_bbox}")

        landmk_5p = np.zeros((5, 2), dtype=np.int)        # 5 keypoints from landmark
        landmk_5p[0] = np.round(np.mean(landmk[36:42,:], axis=0)).astype("int")  # eye left
        landmk_5p[1] = np.round(np.mean(landmk[42:48,:], axis=0)).astype("int")  # eye right 
        landmk_5p[2] = landmk[33]                                                # nose
        landmk_5p[3] = landmk[48]                                                # mouth left
        landmk_5p[4] = landmk[54]                                                # mouth right

        if is_save :
            np.savetxt(str(path_landmk), landmk,    fmt="%d", delimiter=" ", newline="\n", encoding="utf-8")
            np.savetxt(str(path_landmk_5p), landmk_5p, fmt="%d", delimiter=" ", newline="\n", encoding="utf-8")
            np.savetxt(str(path_landmk_bbox), landmk_bbox, fmt="%d", delimiter=" ", newline="\n", encoding="utf-8")
            
        # if show_image:   #for show the image
        #   for (x, y) in landmk:
        #       cv2.circle(image, (x, y), 1, (0, 0, 255), -1)
    
    return True


if __name__ == "__main__":

    path_data = Path("/mnt/data/DATA/GRID/processing")

    v_ids = sorted(path_data.glob("*"))
    print(f"Have {len(v_ids)} dir files, \n {v_ids[:3]}")

    # for each person
    for i, ipath_id in enumerate(v_ids):
        id_name = ipath_id.stem

        dir_video = path_data.joinpath(id_name)
        v_dir_videos = sorted(dir_video.glob("*/video/video_orignal.mp4"))

        print(f"\nid= {id_name}, Have {len(v_dir_videos)} video files, \n {v_dir_videos[:3]}")

        # for each video
        for ipath_video in v_dir_videos:
            
            ipath_data = ipath_video.parent.parent
            #print(f"ipath_data = {ipath_data}")
            
            ipath_frames = ipath_data.joinpath("frames")
            ipath_landmk = ipath_data.joinpath("landmk")
            ipath_landmk_5p = ipath_data.joinpath("landmk_5p")
            ipath_landmk_bbox = ipath_data.joinpath("landmk_bbox")

            ipath_landmk.mkdir(parents=True, exist_ok=True)
            ipath_landmk_5p.mkdir(parents=True, exist_ok=True)
            ipath_landmk_bbox.mkdir(parents=True, exist_ok=True)


            # STEP1: get frames
            v_path_frames = sorted(ipath_frames.glob("*.jpg"))
            print(f"--> Extract landmark: have {len(v_path_frames)} frames in {ipath_frames}")

            # STEP2: extract landmarks
            for k, kpath_frame in enumerate(v_path_frames):
                frame_id = int(kpath_frame.stem[-4:])
                ipath_landmk_frame = ipath_landmk.joinpath(f"frame-{frame_id:04d}.txt")   
                ipath_landmk_frame_5p = ipath_landmk_5p.joinpath(f"frame-{frame_id:04d}.txt")    
                ipath_landmk_frame_bbox = ipath_landmk_bbox.joinpath(f"frame-{frame_id:04d}.txt")


                if not ipath_landmk_frame.exists() or not ipath_landmk_frame_5p.exists() or not ipath_landmk_frame_bbox.exists():
                    
                    # detect image
                    image = cv2.imread(str(kpath_frame)) 
                    detect_image(kpath_frame, 
                                path_landmk=ipath_landmk_frame, 
                                path_landmk_5p=ipath_landmk_frame_5p, 
                                path_landmk_bbox=ipath_landmk_frame_bbox,
                                image=image, is_save=True)

                    # if frame_id % 10 == 0:
                    #     print(f"extract and save the {frame_id}-th frame")
            #break
        #break